import cv2
import os
import numpy as np
from tqdm import tqdm
import multiprocessing
from multiprocessing import Pool
from functools import partial
import matplotlib.pyplot as plt

class ImageCropper:
    def __init__(self):
        self.reference_right_edge = None  # 参考右边缘位置
        
    def find_right_edge(self, binary_img):
        # 计算垂直投影
        projection = np.sum(binary_img, axis=0)
        
        # 使用一维高斯卷积减少噪声影响
        window = 15
        sigma = window/3.0
        x = np.linspace(-window//2, window//2, window)
        gaussian = np.exp(-(x**2)/(2*sigma**2))
        gaussian = gaussian/np.sum(gaussian)
        smoothed = np.convolve(projection, gaussian, mode='valid')
        
        # 计算梯度
        gradient = np.gradient(smoothed)
        
        # 从右向左扫描找到第一个显著边缘
        threshold = np.std(gradient) * 2
        for i in range(len(gradient)-1, window, -1):
            if abs(gradient[i]) > threshold:
                # 使用局部最大值优化边缘位置
                local_region = gradient[max(0, i-5):min(len(gradient), i+6)]
                local_offset = np.argmax(np.abs(local_region))
                return i - 5 + local_offset
        return None
    
    def initialize_with_reference(self, reference_img):
        # 转换为灰度图
        gray = cv2.cvtColor(reference_img, cv2.COLOR_BGR2GRAY)
        
        # 增强对比度
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)
        
        # 检测垂直线
        binary = detect_vertical_lines(gray)
        
        # 找到右边缘
        right_edge = self.find_right_edge(binary)
        if right_edge is not None:
            self.reference_right_edge = right_edge
    
    def crop_image(self, img):
        # 转换为灰度图并增强对比度
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)
        
        # 检测垂直线
        binary = detect_vertical_lines(gray)
        
        # 找到当前帧的右边缘
        current_right_edge = self.find_right_edge(binary)
        
        if current_right_edge is None:
            return img
            
        # 直接切掉右边缘右侧的内容
        cropped = img[:, :current_right_edge]
        
        return cropped

def detect_vertical_lines(gray_img):
    # 使用Sobel算子增强垂直边缘检测
    sobelx = cv2.Sobel(gray_img, cv2.CV_64F, 1, 0, ksize=5)
    sobelx = cv2.convertScaleAbs(sobelx)
    
    # 使用自适应阈值
    binary = cv2.adaptiveThreshold(sobelx, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                 cv2.THRESH_BINARY, 21, 2)
    
    # 形态学操作增强垂直线条
    kernel_vertical = np.ones((15,1), np.uint8)
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel_vertical)
    binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel_vertical)
    
    return binary

def process_single_image(args):
    img_path, cropper = args
    try:
        # 读取当前帧
        current_frame = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        
        if current_frame is None:
            return None
        
        # 裁剪图像
        cropped_frame = cropper.crop_image(current_frame)
        
        return cropped_frame
    except Exception as e:
        print(f"处理图片 {img_path} 时发生错误: {str(e)}")
        return None

def process_images_to_video(processed_frames, output_path, fps=60.0):
    """
    将处理后的帧生成视频
    
    Args:
        processed_frames: 已处理的帧列表
        output_path: 输出视频路径
        fps: 视频帧率，默认30fps
    """
    try:
        if not processed_frames:
            print("错误：没有可用的处理帧")
            return
            
        # 获取视频尺寸
        height, width = processed_frames[0].shape[:2]
        
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # 创建视频写入器
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        expected_duration = len(processed_frames) / fps
        print(f"预期视频时长: {expected_duration:.2f} 秒")
        
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        
        # 写入所有帧
        for frame in tqdm(processed_frames, desc="生成视频"):
            out.write(frame)
        
        out.release()
        print(f"视频生成完成！实际时长: {expected_duration:.2f} 秒")
        
    except Exception as e:
        print(f"生成视频时发生错误: {str(e)}")

def process_images(input_path):
    """
    处理输入文件夹中的所有图片
    
    Args:
        input_path: 输入图片文件夹路径
    
    Returns:
        处理后的帧列表
    """
    try:
        # 获取所有bmp文件
        image_files = [f for f in os.listdir(input_path) if f.lower().endswith('.bmp')]
        if not image_files:
            print(f"错误：在 {input_path} 中没有找到BMP文件")
            return []
        image_files.sort()
        
        total_images = len(image_files)
        print(f"总共找到 {total_images} 张图片")
        
        # 读取第一张图片作为参考帧
        first_image_path = os.path.join(input_path, image_files[0])
        print(f"正在读取第一张图片: {first_image_path}")
        reference_frame = cv2.imdecode(np.fromfile(first_image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        
        if reference_frame is None:
            print(f"错误：无法读取图片 {first_image_path}")
            return []
        
        # 创建裁剪器并初始化
        cropper = ImageCropper()
        cropper.initialize_with_reference(reference_frame)
        
        # 准备并行处理的参数
        process_args = [(os.path.join(input_path, img_file), cropper) for img_file in image_files]
        
        # 获取CPU核心数
        num_cores = multiprocessing.cpu_count()
        print(f"使用 {num_cores} 个CPU核心进行并行处理")
        
        # 创建进程池
        with Pool(num_cores) as pool:
            # 使用tqdm显示进度
            processed_frames = list(tqdm(
                pool.imap(process_single_image, process_args),
                total=len(process_args),
                desc="处理图片"
            ))
        
        # 过滤掉处理失败的帧
        processed_frames = [frame for frame in processed_frames if frame is not None]
        
        if not processed_frames:
            print("错误：没有成功处理任何图片")
            return []
            
        print(f"成功处理 {len(processed_frames)} 张图片")
        return processed_frames
        
    except Exception as e:
        print(f"处理图片时发生错误: {str(e)}")
        return []

def analyze_edge_movement(input_path):
    """分析边缘移动并生成位移图"""
    # 获取所有bmp文件
    image_files = [f for f in os.listdir(input_path) if f.lower().endswith('.bmp')]
    if not image_files:
        print(f"错误：在 {input_path} 中没有找到BMP文件")
        return
    image_files.sort()
    
    # 创建裁剪器
    cropper = ImageCropper()
    
    # 读取第一帧并初始化
    first_image_path = os.path.join(input_path, image_files[0])
    reference_frame = cv2.imdecode(np.fromfile(first_image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    cropper.initialize_with_reference(reference_frame)
    
    # 存储每一帧的边缘位置
    edge_positions = []
    frame_numbers = []
    
    # 处理每一帧
    for i, img_file in enumerate(tqdm(image_files, desc="分析边缘位移")):
        img_path = os.path.join(input_path, img_file)
        frame = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        
        # 转换为灰度图并增强对比度
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)
        
        # 检测垂直线
        binary = detect_vertical_lines(gray)
        
        # 找到当前帧的右边缘
        current_edge = cropper.find_right_edge(binary)
        
        if current_edge is not None:
            edge_positions.append(current_edge)
            frame_numbers.append(i)
    
    # 计算相对位移（相对于第一帧）
    edge_displacements = [pos - edge_positions[0] for pos in edge_positions]
    
    # 绘制位移图
    plt.figure(figsize=(12, 6))
    plt.plot(frame_numbers, edge_displacements, 'b-', label='边缘位移')
    plt.xlabel('帧数')
    plt.ylabel('位移 (像素)')
    plt.title('边缘位移随时间变化')
    plt.grid(True)
    plt.legend()
    
    # 保存图表
    plt.savefig(os.path.join(os.path.dirname(input_path), 'edge_displacement.png'))
    plt.close()
    
    # 输出统计信息
    print(f"总帧数: {len(frame_numbers)}")
    print(f"最大位移: {max(edge_displacements):.2f} 像素")
    print(f"最小位移: {min(edge_displacements):.2f} 像素")
    print(f"平均位移: {np.mean(edge_displacements):.2f} 像素")
    print(f"位移标准差: {np.std(edge_displacements):.2f} 像素")

def visualize_edge_detection(input_path, output_folder="edge_visualization"):
    """可视化边缘检测结果"""
    # 创建输出文件夹
    os.makedirs(output_folder, exist_ok=True)
    
    # 获取所有bmp文件
    image_files = [f for f in os.listdir(input_path) if f.lower().endswith('.bmp')]
    if not image_files:
        print(f"错误：在 {input_path} 中没有找到BMP文件")
        return
    image_files.sort()
    
    # 创建裁剪器
    cropper = ImageCropper()
    
    # 读取第一帧并初始化
    first_image_path = os.path.join(input_path, image_files[0])
    reference_frame = cv2.imdecode(np.fromfile(first_image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    cropper.initialize_with_reference(reference_frame)
    
    # 处理每一帧
    for i, img_file in enumerate(tqdm(image_files[:100], desc="生成边缘可视化图像")):  # 默认只处理前100帧
        img_path = os.path.join(input_path, img_file)
        frame = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        
        # 转换为灰度图并增强对比度
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)
        
        # 检测垂直线
        binary = detect_vertical_lines(gray)
        
        # 找到当前帧的右边缘
        current_edge = cropper.find_right_edge(binary)
        
        if current_edge is not None:
            # 创建可视化图像
            vis_img = frame.copy()
            
            # 在原图上绘制检测到的边缘线
            cv2.line(vis_img, (current_edge, 0), (current_edge, vis_img.shape[0]), (0, 255, 0), 2)
            
            # 添加文本信息
            cv2.putText(vis_img, f"Frame: {i}", (10, 30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            cv2.putText(vis_img, f"Edge: {current_edge}", (10, 70), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            
            # 创建组合图像
            binary_colored = cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR)
            combined = np.hstack((vis_img, binary_colored))
            
            # 保存图像
            output_path = os.path.join(output_folder, f"edge_frame_{i:04d}.jpg")
            cv2.imwrite(output_path, combined)
    
    print(f"边缘检测可视化完成！结果保存在 {output_folder} 文件夹中")

def analyze_image_dimensions(input_path):
    """分析每张图片裁剪后的尺寸并生成可视化图表"""
    # 获取所有bmp文件
    image_files = [f for f in os.listdir(input_path) if f.lower().endswith('.bmp')]
    if not image_files:
        print(f"错误：在 {input_path} 中没有找到BMP文件")
        return
    image_files.sort()
    
    # 创建裁剪器
    cropper = ImageCropper()
    
    # 读取第一帧并初始化
    first_image_path = os.path.join(input_path, image_files[0])
    reference_frame = cv2.imdecode(np.fromfile(first_image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    cropper.initialize_with_reference(reference_frame)
    
    # 存储尺寸数据
    frame_numbers = []
    widths = []
    heights = []
    
    # 处理每一帧
    for i, img_file in enumerate(tqdm(image_files, desc="分析图片尺寸")):
        img_path = os.path.join(input_path, img_file)
        frame = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        
        # 裁剪图像
        cropped_frame = cropper.crop_image(frame)
        
        # 记录尺寸
        height, width = cropped_frame.shape[:2]
        frame_numbers.append(i)
        widths.append(width)
        heights.append(height)
    
    # 创建图表
    plt.figure(figsize=(15, 10))
    
    # 绘制宽度变化
    plt.subplot(2, 1, 1)
    plt.plot(frame_numbers, widths, 'b-', label='宽度')
    plt.xlabel('帧数')
    plt.ylabel('宽度 (像素)')
    plt.title('裁剪后图片宽度随时间变化')
    plt.grid(True)
    plt.legend()
    
    # 绘制高度变化
    plt.subplot(2, 1, 2)
    plt.plot(frame_numbers, heights, 'r-', label='高度')
    plt.xlabel('帧数')
    plt.ylabel('高度 (像素)')
    plt.title('裁剪后图片高度随时间变化')
    plt.grid(True)
    plt.legend()
    
    # 调整子图之间的间距
    plt.tight_layout()
    
    # 保存图表
    plt.savefig(os.path.join(os.path.dirname(input_path), 'image_dimensions.png'))
    plt.close()
    
    # 输出统计信息
    print(f"\n尺寸统计信息:")
    print(f"宽度: 最小={min(widths)}px, 最大={max(widths)}px, 平均={np.mean(widths):.2f}px, 标准差={np.std(widths):.2f}px")
    print(f"高度: 最小={min(heights)}px, 最大={max(heights)}px, 平均={np.mean(heights):.2f}px, 标准差={np.std(heights):.2f}px")
    print(f"宽度变化范围: {max(widths) - min(widths)}px")

if __name__ == '__main__':
    input_path = r"F:\边缘test"  # 您的输入路径
    output_path = r"F:\test\output_video_right_cropped.mp4"  # 新的输出路径
    
    # 分析图片尺寸
    analyze_image_dimensions(input_path)
    
    # 处理所有图片
    processed_frames = process_images(input_path)
    
    # 生成视频
    if processed_frames:
        process_images_to_video(processed_frames, output_path)
